import cv2
import numpy as np
import torch
from collections import defaultdict

from htm.Utils import Utils


def runExperiment():
    np.random.seed(1)  # The only place to set seed
    imgPath = "../visionBodyLearn/outputImages/resaved_temp1.png"
    origScene2 = cv2.imread(imgPath)
    edgeScene2 = Utils.getEdgeImg(origScene2)
    scene2 = Utils.encodeScenarioUsingImage(origScene2, edgeScene2)
    l46Exp = torch.load('trained_models/piunExp.m')
    edgeScene1 = torch.load('trained_models/edgeScene1.m')
    changesInScene = edgeScene2.astype(int) - edgeScene1.astype(int)
    posEdge = np.zeros_like(edgeScene1)
    posEdge[np.where(changesInScene == 1)] = 1
    negEdge = np.zeros_like(edgeScene1)
    negEdge[np.where(changesInScene == -1)] = 1
    posEdgeForHtm = Utils.encodeScenarioUsingImage(origScene2, posEdge)
    negEdgeForHtm = Utils.encodeScenarioUsingImage(origScene2, negEdge)
    convergence = defaultdict(int)
    for objectDescription in [scene2]:
        print('inferring Object:', objectDescription['name'])
        numSensationsToInference = l46Exp.inferObjectSimple(objectDescription)
        convergence[numSensationsToInference] += 1
    sortedKeys = sorted([key for key in convergence.keys()])
    for key in sortedKeys:
        print('infer steps:', key, 'infered object num:', convergence[key])


if __name__ == "__main__":
    runExperiment()
